Description

This PR adds a native-QAT path for the routed-expert GMMs, on top of the existing fake QAT (fp8_moe_fake_quant + fp8_fake_quant_hoist=gather8).

What fake QAT does today:

• It all-gathers the FP8 weight codes.
• It then dequantizes them to a BF16 copy in HBM.
• Every GMM runs in BF16.

What the new flags change:

• fp8_moe_native_fwd=true: the wi/wo forward GMMs run FP8×FP8 on the TPU7x MXU. They read the gathered codes and per-channel scale directly, and quantize the activation inside the kernel.
• fp8_moe_native_bwd_dlhs=true: the backward dLHS GMM also reads the codes and dequantizes each tile in VMEM. No BF16 copy of the weight is written to HBM.
• Numerics: the backward is bit-identical to the existing dequant path. The forward changes numerics only by quantizing activations to FP8.

All new flags default to off. With them off, the MoE GMM programs are byte-identical to base.

Related: the dropless full-buffer index clamp that was part of this PR is now #6068. The end-to-end results below were measured with that clamp in the tree; land #6068 first if you run dropless replay with moe_routing_index_math=count.

Why:

• The trainer already holds FP8 weight codes after the gather8 all-gather. Fake QAT spends HBM bandwidth and time writing BF16 copies, then runs BF16 GMMs at half the TPU7x FP8 MXU throughput.
• Running the forward in FP8 also matches what the FP8 rollout's MoE kernel computes.
• On the Qwen3.5-397B RL trainer bench (128 chips, same image and HBM cap per pair), step time drops 2.56–3.03% vs fake QAT (throughput +2.63–3.13%). The range depends on the routing setting (see Tests).
• In a GMM kernel microbenchmark on one TPU7x TensorCore at the trainer's per-device shapes, the native forward kernels (FP8 GMM, including the in-kernel BF16 → FP8 activation quantization) are 1.6–1.7× faster than the BF16 GMM, which needs no activation quantization (wi 1.635×, wo 1.731×).
• At one seed, convergence shows no regression signal.

Data flow

flowchart LR
  subgraph F["Fake QAT"]
    F1["FP8 weights<br/>+ scales"] --> F2["dequant<br/>to BF16"]
    F2 --> F3["fwd GMM<br/>BF16 x BF16"]
    F2 --> F4["bwd dLHS<br/>BF16 x BF16"]
    F2 --> F5["bwd dW<br/>BF16"]
  end
  subgraph N["Native QAT (this PR)"]
    N1["FP8 weights<br/>+ scales"] --> N3["fwd GMM<br/>FP8 x FP8"]
    N1 --> N4["bwd dLHS<br/>VMEM dequant"]
    N1 --> N2["dequant<br/>to BF16"] --> N5["bwd dW<br/>BF16"]
  end

Gradient anchor: an FP8 primal would get an FP8 cotangent. So each weight also enters the op as its BF16 dequantized value, an extra input whose value the forward never reads; XLA drops that unread dequantization. The weight gradient is returned as that input's cotangent. It then flows through dequantize_weight_codes_ste to the FSDP shard exactly as it does today.

Implementation

• kernels/megablox/ops.py adds four functions:

  • gmm_silu_fp8: wi with fused SwiGLU.
  • gmm_dswiglu_fp8: wo.
  • gmm_fp8_forward: forward only, for ring_ragged_unsort_down_projection.
  • dequantize_fp8_codes.

  Each custom_vjp forward calls gmm_v2(rhs=codes, rhs_scale=…, maybe_quantize_lhs=True). With fp8_moe_native_bwd_dlhs, the dLHS backward calls gmm_v2 with a codes RHS (rhs_contract_scale). Otherwise the backward dequantizes and reuses the existing BF16 backward.

• pallas_mosaic_tpu_v2_gmm_kernel.py:

  • LHS quantization options: lhs_quant_block_size (512 or full K), lhs_f32_scale (max(amax, 1e-12)/448 in f32), lhs_hoist_quant and rhs_scale_first. fp8_moe_native_fwd always uses full K with an f32 scale and f32 accumulation; the other options are ops-level arguments for benchmarks and tests only.
  • The rhs_contract_dequant path: an FP8 double buffer with per-tile dequant fused into each matmul. With one K tile (production), the dequantized block is also cached in VMEM across its m tiles.
  • Tiles that path can't take (tile_n % 32 != 0, tile_k % 128 != 0, non-dividing tiles) fall back to the XLA dequant. Production tiles take the native path.
  • The VMEM cache is allocated only with one K tile, and only its fill / cached pair is dispatched. With more K tiles there is no cache tile: each step dequantizes straight into its matmul; outputs are unchanged and it is still faster than the XLA-dequant fallback (see Tests).

• layers/moe.py:

  • Adds an Fp8NativeWeight(anchor, codes, scale) pytree, built from the gather8 codes after the prefetch barrier. The shard-map in_specs give the scale P(e, None, n).
  • Dispatch points:
    • gmm_up_or_silu (wi);
    • _run_wo_gmm_dswiglu (wo);
    • down_projection / down_projection_bwd_dswiglu inside _ring_down_project_and_unpermute (ring-of-experts wo).
  • A RoutedMoE call with fp8_moe_native_fwd=true but no gather8 absmax now raises instead of silently running fake QAT. This happens for unscanned layers, KV-cache decode, or a Qwix-intercepted call.
  • RoutedMoE also raises if the traced w0/wo partition specs split K (for example through custom_mesh_and_rule or parallelism -1), since the activation scale must span the whole K.

• configs/types.py, configs/base.yml: two new flags, both default off.

  Flag	Values	Recommended
  fp8_moe_native_fwd	bool	true
  fp8_moe_native_bwd_dlhs	bool	true

  • fp8_moe_native_fwd requires fp8_moe_fake_quant, fp8_fake_quant_hoist=gather8, sparse_matmul, use_tokamax_gmm, use_gmm_v2, moe_wi_gmm_fused_swiglu, moe_wo_gmm_fused_dswiglu, fixed tiles (use_gmm_v2_heuristic_tiling=false: the heuristic tiling budgets a bf16 accumulator when the LHS is quantized in the kernel, but the native forward accumulates in f32), and single-K-tile forward tiles (wi_tile_fwd_embed_dim = embed dim and wo_tile_fwd_mlp_dim = MoE MLP dim, plus eval_* when separate eval axis rules are set; make_gmm_configs also rejects tile_k != size_k when lhs_quant_block_size=0), and an unsplit K in both forward GMMs (no mesh axis with parallelism > 1 on wi's embed dim or wo's MoE MLP dim, and num_moe_emb_chunks=0), since the activation scale spans the whole K.
  • fp8_moe_native_bwd_dlhs additionally requires fp8_moe_native_fwd and moe_gmm_v2_dlhs_transpose_rhs.

Follow-ups:

• FP8 backward for the weight-gradient TGMM (dW), which is still BF16.
• Make fp8_moe_native_fwd imply fp8_fake_quant_hoist=gather8 and the fake-QAT flags; today it requires them, because it reuses gather8's FP8 weights.

Tests

Unit and correctness

Performance: Qwen3.5-397B-A17B, DeepSWE RL, TPU7x. Fake = base (caafa5bc6). Native = base + fp8_moe_native_fwd=true fp8_moe_native_bwd_dlhs=true, with full-K f32 activation scales and f32 accumulation (set then by two flags that are now the fixed behavior). Both arms use the same image and --xla_tpu_max_hbm_size_mib=90112.

Grouped matmul (GMM) kernel microbenchmark (one TPU7x TensorCore, the trainer's per-device shapes and production tiles):

• Setting: kernel only, one call at full fill (64 experts, 225,280 rows).
  • The BF16 arm gets a weight that was dequantized to BF16 beforehand, so fake QAT's XLA weight dequant is not timed.
  • The FP8 arm takes the FP8 weight and scale directly, and its timing includes the in-kernel activation quantization.
• Conclusion: the FP8 kernels are 1.6–1.7× faster even after paying for activation quantization. Against the full fake-QAT path, which also dequantizes the weight in XLA, the forward gain is larger.

Trainer bench (128 chips, 4x4x8, mean of steps 1–4; tok/s/chip = 37,748,736 tokens per step ÷ (128 × s/step)):

Row 2 overflows the expert buffer on one micro-batch per step, and the dropless path reruns it (+16 s/step in both arms). That rerun needs the index clamp from #6068, which these runs had.

640-chip e2e (seed 677013886):

XProf (one trainer micro-batch, TPU:0):

Total device busy time isn't compared here: the two runs sampled different micro-batches (attention time alone differs by 247 ms). The step-time gain is measured in the trainer bench above.

Convergence:

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.